package com.xy.tool.concurrent;

import lombok.extern.slf4j.Slf4j;
import org.springframework.util.CollectionUtils;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.Random;
import java.util.concurrent.*;
import java.util.function.Consumer;

/**
 * 批量任务执行
 *
 * @param <T>
 */
@Slf4j
public class BatchTaskExecutor<T> {

    private final static int DEFAULT_SIZE_THRESHOLD = 100;
    private final static int DEFAULT_TIME_THRESHOLD = 5000;
    private final static int DEFAULT_POLL_TIME_OUT = 1000;
    private final static int DEFAULT_QUEUE_SIZE = 1024;

    private int sizeThreshold = DEFAULT_SIZE_THRESHOLD;
    private int timeThreshold = DEFAULT_TIME_THRESHOLD;
    private int pollTimeOut = DEFAULT_POLL_TIME_OUT;
    private int queueSize = DEFAULT_QUEUE_SIZE;
    private volatile boolean running;

    private ExecutorService mainThreadGroup;
    private ExecutorService workThreadGroup;
    private BlockingQueue<T> queue;

    private final Consumer<List<T>> batchConsumeDelegate;

    /**
     * 构建方法
     *
     * @param batchConsumeDelegate      批量消费代理
     */
    public BatchTaskExecutor(Consumer<List<T>> batchConsumeDelegate) {
        Objects.requireNonNull(batchConsumeDelegate);
        this.batchConsumeDelegate = batchConsumeDelegate;

        queue = new LinkedBlockingQueue<>(queueSize);
    }

    /**
     * 设置批量消费大小的阈值
     *
     * @param sizeThreshold
     */
    public void setSizeThreshold(int sizeThreshold) {
        if(sizeThreshold < 1) {
            throw new IllegalArgumentException();
        }
        this.sizeThreshold = sizeThreshold;
    }

    /**
     * 设置批量消费时间的阈值
     *
     * @param timeThreshold
     */
    public void setTimeThreshold(int timeThreshold) {
        if(timeThreshold < 1) {
            throw new IllegalArgumentException();
        }
        this.timeThreshold = timeThreshold;
    }

    /**
     * 设置轮询任务的超时时间
     *
     * @param pollTimeOut
     */
    public void setPollTimeOut(int pollTimeOut) {
        if(pollTimeOut < 1) {
            throw new IllegalArgumentException();
        }
        this.pollTimeOut = pollTimeOut;
    }

    /**
     * 设置任务队列大小
     *
     * @param queueSize
     */
    public void setQueueSize(int queueSize) {
        if(queueSize < 1) {
            throw new IllegalArgumentException();
        }
        this.queueSize = queueSize;
    }

    /**
     * 设置任务工作线程池
     *
     * @param workThreadGroup
     */
    public void setWorkThreadGroup(ExecutorService workThreadGroup) {
        Objects.requireNonNull(workThreadGroup);
        this.workThreadGroup = workThreadGroup;
    }

    /**
     * 启动
     */
    public void start() {
        if(!running) {
            synchronized (this) {
                if(!running) {
                    running = true;
                    start0();
                }
            }
        }
    }

    private void start0() {
        mainThreadGroup = Executors.newSingleThreadExecutor();
        mainThreadGroup.execute(() -> {
            int size = 0;
            long time = 0, timeStart;
            T current;
            BlockingQueue<T> queue0 = queue;
            List<T> tempList = new ArrayList<>();
            for(;;) {
                timeStart = System.currentTimeMillis();
                try {
                    current = queue0.poll(pollTimeOut, TimeUnit.MILLISECONDS);
                } catch (InterruptedException e) {
                    break;
                }
                time += (System.currentTimeMillis() - timeStart);
                if(current != null) {
                    tempList.add(current);
                    size++;
                } else{
                    if(!running) {
                        break;
                    }
                }
                if(size > sizeThreshold
                        || (size > 0 && time > timeThreshold)) {
                    batchConsume(tempList);
                    size = 0;
                    time = 0;
                    tempList.clear();
                }
            }
            //退出主循环后的收尾处理
            batchConsume(tempList);
        });
    }

    private void batchConsume(List<T> tempList) {
        if(!CollectionUtils.isEmpty(tempList)) {
            if(workThreadGroup != null) {
                List<T> list = new ArrayList<>(tempList);
                workThreadGroup.submit(() -> {
                    try {
                        batchConsumeDelegate.accept(list);
                    } catch (Exception e) {
                        log.error("Exception occurs when batch consume.", e);
                    }
                });
            } else {
                try {
                    batchConsumeDelegate.accept(tempList);
                } catch (Exception e) {
                    log.error("Exception occurs when batch consume.", e);
                }
            }
        }
    }

    /**
     * 停止
     */
    public void stop() {
        if(running) {
            synchronized (this) {
                if(running) {
                    running = false;
                    stop0();
                }
            }
        }
    }

    private void stop0() {
        //停止的一瞬间，队列中仍有可能放入新数据，收尾处理。
        List<T> tempList = new ArrayList<>();
        T current;
        for(;;) {
            try {
                current = queue.poll(1, TimeUnit.MILLISECONDS);
            } catch (InterruptedException e) {
                break;
            }
            if(current == null) {
                break;
            }
            tempList.add(current);
        }
        batchConsume(tempList);

        mainThreadGroup.shutdown();
        try {
            // Wait a while for existing tasks to terminate
            if (!mainThreadGroup.awaitTermination(30, TimeUnit.SECONDS)) {
                mainThreadGroup.shutdownNow();
                // Wait a while for tasks to respond to being cancelled
                if (!mainThreadGroup.awaitTermination(30, TimeUnit.SECONDS)) {
                    log.error("Thread pool did not terminate.");
                }
            }
        } catch (InterruptedException ie) {
            // (Re-)Cancel if current thread also interrupted
            mainThreadGroup.shutdownNow();
            // Preserve interrupt status
            Thread.currentThread().interrupt();
        }
    }

    /**
     * 添加任务
     *
     * @param task
     * @return
     */
    public boolean add(T task) {
        if(!running) {
            return false;
        }
        return queue.add(task);
    }

    /**
     * 添加任务
     *
     * @param task
     * @throws InterruptedException
     */
    public void put(T task) throws InterruptedException {
        if(!running) {
            throw new IllegalStateException();
        }
        queue.put(task);
    }

    /**
     * 添加任务
     *
     * @param task
     * @return
     */
    public boolean offer(T task) {
        if(!running) {
            return false;
        }
        return queue.offer(task);
    }

    /**
     * 添加任务
     *
     * @param task
     * @param timeout
     * @param unit
     * @return
     * @throws InterruptedException
     */
    public boolean offer(T task, long timeout, TimeUnit unit)
            throws InterruptedException {
        if(!running) {
            return false;
        }
        return queue.offer(task, timeout, unit);
    }


    /// tests

    public static void main(String[] args) throws IOException, InterruptedException {
        Random r = new Random();
        BatchTaskExecutor<String> batchTaskExecutor = new BatchTaskExecutor<>(
                e -> {
                    for(String s : e) {
                        System.out.println(Thread.currentThread().getId() + ": " + s);
                    }
                });
        batchTaskExecutor.setWorkThreadGroup(Executors.newFixedThreadPool(4));
        batchTaskExecutor.start();
        for(int i = 1; i <= 201; i++) {
            batchTaskExecutor.add("Hello Word: " + i);
            Thread.sleep(r.nextInt(100));
        }
        batchTaskExecutor.stop();
        System.out.println(Thread.currentThread().getId() + " exit");
    }

}
